from flask import Blueprint, request, jsonify, render_template
import pymysql
import numpy as np
from pymysql.cursors import DictCursor
import joblib
import os
from datetime import datetime

# 创建蓝图
predict_bp = Blueprint('predict', __name__, url_prefix='/predict')

# 数据库配置
DB_CONFIG = {
    'host': 'localhost',
    'user': 'root',
    'password': '123456',
    'database': 'Ab_data',
    'charset': 'utf8mb4',
    'cursorclass': DictCursor
}

# 模型保存路径
MODEL_DIR = 'models'
os.makedirs(MODEL_DIR, exist_ok=True)


def get_db_connection():
    """创建数据库连接"""
    return pymysql.connect(**DB_CONFIG)


def encode_sex(sex):
    """性别编码"""
    sex_mapping = {'M': 1, 'F': 0, 'I': -1}
    return sex_mapping.get(sex, -1)


def load_model(model_type):
    """加载已训练的模型"""
    model_path = os.path.join(MODEL_DIR, f'abalone_{model_type}.pkl')
    if os.path.exists(model_path):
        return joblib.load(model_path)
    return None


def update_predictions(abalone_id, predictions):
    """更新数据库中的预测结果"""
    query = """
    UPDATE new_abalone
    SET predicted_rings_LR = %s,
        predicted_rings_DT_R = %s,
        predicted_rings_RF_R = %s,
        prediction_time_LR = NOW(),
        prediction_time_DT_R = NOW(),
        prediction_time_RF_R = NOW()
    WHERE id = %s
    """

    # 四舍五入预测结果
    rounded_predictions = [int(round(p)) for p in predictions]

    connection = get_db_connection()
    try:
        with connection.cursor() as cursor:
            cursor.execute(query, (*rounded_predictions, abalone_id))
        connection.commit()
    finally:
        connection.close()


def get_abalone_features(abalone_id):
    """获取单个鲍鱼的特征数据"""
    query = """
    SELECT sex, length, diameter, height, 
           whole_weight, shucked_weight, 
           viscera_weight, shell_weight
    FROM new_abalone
    WHERE id = %s
    """
    connection = get_db_connection()
    try:
        with connection.cursor() as cursor:
            cursor.execute(query, (abalone_id,))
            result = cursor.fetchone()
            if result:
                return [
                    encode_sex(result['sex']),
                    float(result['length']),
                    float(result['diameter']),
                    float(result['height']),
                    float(result['whole_weight']),
                    float(result['shucked_weight']),
                    float(result['viscera_weight']),
                    float(result['shell_weight'])
                ]
            return None
    finally:
        connection.close()


def get_all_abalones():
    """获取所有鲍鱼数据"""
    query = """
    SELECT id, sex, length, diameter, height, 
           whole_weight, shucked_weight, 
           viscera_weight, shell_weight
    FROM new_abalone
    """
    connection = get_db_connection()
    try:
        with connection.cursor() as cursor:
            cursor.execute(query)
            return cursor.fetchall()
    finally:
        connection.close()


@predict_bp.route('/', methods=['GET', 'POST'])
def predict():
    """单个预测功能"""
    if request.method == 'GET':
        return render_template('predict.html')

    data = request.form
    abalone_id = data.get('abalone_id')

    if not abalone_id:
        return jsonify({'success': False, 'message': '请输入鲍鱼编号'})

    try:
        abalone_id = int(abalone_id)
        features = get_abalone_features(abalone_id)

        if features is None:
            return jsonify({'success': False, 'message': f'未找到编号为 {abalone_id} 的鲍鱼记录'})

        # 加载所有模型
        models = {
            'LR': load_model('LR'),
            'DT-R': load_model('DT-R'),
            'RF-R': load_model('RF-R')
        }

        # 检查所有模型是否存在
        for model_type, model in models.items():
            if model is None:
                return jsonify({'success': False, 'message': f'找不到 {model_type} 模型，请确保模型已存在'})

        # 使用所有模型进行预测
        predictions = {
            model_type: int(round(model.predict([features])[0]))
            for model_type, model in models.items()
        }

        # 更新数据库
        update_predictions(abalone_id, list(predictions.values()))

        return jsonify({
            'success': True,
            'predictions': predictions,
            'abalone_id': abalone_id
        })
    except ValueError:
        return jsonify({'success': False, 'message': '请输入有效的数字编号'})
    except Exception as e:
        return jsonify({'success': False, 'message': str(e)})


@predict_bp.route('/batch', methods=['POST'])
def batch_predict():
    """批量预测功能"""
    model_type = request.form.get('model_type', 'LR')

    try:
        model = load_model(model_type)
        if model is None:
            return jsonify({'success': False, 'message': f'找不到 {model_type} 模型，请确保模型已存在'})

        abalones = get_all_abalones()
        if not abalones:
            return jsonify({'success': True, 'message': '没有可预测的鲍鱼数据', 'count': 0})

        # 准备特征数据
        features_list = []
        id_list = []
        for abalone in abalones:
            features = [
                encode_sex(abalone['sex']),
                float(abalone['length']),
                float(abalone['diameter']),
                float(abalone['height']),
                float(abalone['whole_weight']),
                float(abalone['shucked_weight']),
                float(abalone['viscera_weight']),
                float(abalone['shell_weight'])
            ]
            features_list.append(features)
            id_list.append(abalone['id'])

        # 批量预测
        predictions = model.predict(features_list)

        # 更新数据库
        connection = get_db_connection()
        try:
            with connection.cursor() as cursor:
                for abalone_id, prediction in zip(id_list, predictions):
                    # 只更新当前模型的预测结果
                    if model_type == 'LR':
                        field = 'predicted_rings_LR'
                        time_field = 'prediction_time_LR'
                    elif model_type == 'DT-R':
                        field = 'predicted_rings_DT_R'
                        time_field = 'prediction_time_DT_R'
                    elif model_type == 'RF-R':
                        field = 'predicted_rings_RF_R'
                        time_field = 'prediction_time_RF_R'
                    else:
                        continue

                    query = f"""
                    UPDATE new_abalone
                    SET {field} = %s,
                        {time_field} = NOW()
                    WHERE id = %s
                    """
                    cursor.execute(query, (int(round(prediction)), abalone_id))
            connection.commit()
        finally:
            connection.close()

        return jsonify({
            'success': True,
            'message': f'成功使用 {model_type} 模型预测了 {len(abalones)} 条鲍鱼数据',
            'count': len(abalones),
            'model_type': model_type
        })
    except Exception as e:
        return jsonify({'success': False, 'message': str(e)})


@predict_bp.route('/predict_all_models', methods=['POST'])
def predict_all_models():
    """使用所有模型进行预测"""
    try:
        abalones = get_all_abalones()
        if not abalones:
            return jsonify({'success': True, 'message': '没有可预测的鲍鱼数据', 'count': 0})

        # 准备特征数据
        features_list = []
        id_list = []
        for abalone in abalones:
            features = [
                encode_sex(abalone['sex']),
                float(abalone['length']),
                float(abalone['diameter']),
                float(abalone['height']),
                float(abalone['whole_weight']),
                float(abalone['shucked_weight']),
                float(abalone['viscera_weight']),
                float(abalone['shell_weight'])
            ]
            features_list.append(features)
            id_list.append(abalone['id'])

        # 加载所有模型
        models = {
            'LR': load_model('LR'),
            'DT-R': load_model('DT-R'),
            'RF-R': load_model('RF-R')
        }

        # 检查所有模型是否存在
        for model_type, model in models.items():
            if model is None:
                return jsonify({'success': False, 'message': f'找不到 {model_type} 模型，请确保模型已存在'})

        # 批量预测并更新数据库
        connection = get_db_connection()
        try:
            with connection.cursor() as cursor:
                for abalone_id, features in zip(id_list, features_list):
                    predictions = [
                        int(round(models['LR'].predict([features])[0])),
                        int(round(models['DT-R'].predict([features])[0])),
                        int(round(models['RF-R'].predict([features])[0]))
                    ]
                    update_predictions(abalone_id, predictions)
            connection.commit()
        finally:
            connection.close()

        return jsonify({
            'success': True,
            'message': f'成功使用所有模型预测了 {len(abalones)} 条鲍鱼数据',
            'count': len(abalones)
        })
    except Exception as e:
        return jsonify({'success': False, 'message': str(e)})